package option

import (
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"

	"gitee.com/hasika/common-server/codec"
	"gitee.com/hasika/common-server/conn/inet"
	"gitee.com/hasika/gotool/errorx"
)

type MessageHandler func(data codec.Packet, conn inet.Conn) (errorx.IError, proto.Message)

type MessageRegInfo struct {
	Handler      MessageHandler
	ReplyType    inet.ConnectionType
	NilMustReply bool
}

type ProtoMessageRouter struct {
	registry map[protoreflect.MessageType]*MessageRegInfo
}

func NewProtoMessageRouter() *ProtoMessageRouter {
	return &ProtoMessageRouter{
		registry: make(map[protoreflect.MessageType]*MessageRegInfo),
	}
}

func (c *ProtoMessageRouter) Reg(k protoreflect.MessageType, info *MessageRegInfo) {
	c.registry[k] = info
}

func (c *ProtoMessageRouter) RegSimple(proto proto.Message, Handler MessageHandler, ReplyType inet.ConnectionType, NilMustReply ...bool) {
	nilMustReply := false
	if len(NilMustReply) == 1 {
		nilMustReply = NilMustReply[0]

	}
	reqType := proto.ProtoReflect().Type()
	c.registry[reqType] = &MessageRegInfo{
		Handler:      Handler,
		ReplyType:    ReplyType,
		NilMustReply: nilMustReply,
	}
}

func (c *ProtoMessageRouter) Get(data codec.Packet) *MessageRegInfo {
	return c.registry[data.GetMessage().ProtoReflect().Type()]
}

func (c *ProtoMessageRouter) Merge(router *ProtoMessageRouter) {
	for k, v := range router.registry {
		c.registry[k] = v
	}
}

func (c *ProtoMessageRouter) Range(f func(k protoreflect.MessageType, v *MessageRegInfo)) {
	for k, v := range c.registry {
		f(k, v)
	}
}
